import torch
import torch.nn as nn
from transformers import PretrainedConfig
import math

"""
原代码__init__和forward含大量 NPU / 分布式逻辑，此 Demo 简化为Windows 单卡环境，用 PyTorch 原生层替换自定义并行层，保留核心结构（LoRA 投影、RoPE、注意力计算）。
移除 NPU 依赖：删除torch_npu/torchair相关代码，用 PyTorch 原生einsum实现注意力计算；
并行层模拟：ReplicatedLinear/ColumnParallelLinear简化为nn.Linear，适配单卡环境；
数据类型：用float32替代bfloat16，避免 Windows CPU 对bfloat16的兼容性问题；
RoPE 实现：用基础旋转公式替换npu_interleave_rope，确保 Windows CPU/GPU 可运行。
"""

def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    """原代码YARN缩放因子计算函数：补偿长序列RoPE的注意力衰减"""
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


# -------------------------- 模拟原代码中缺失的自定义层（Windows适配） --------------------------
class ReplicatedLinear(nn.Linear):
    """模拟omni的ReplicatedLinear（Windows单卡环境下等同于nn.Linear）"""

    def forward(self, x):
        # 适配原代码返回格式：(输出张量, None)
        return super().forward(x), None


class ColumnParallelLinear(nn.Linear):
    """模拟vllm的ColumnParallelLinear（Windows单卡环境下等同于nn.Linear）"""

    def forward(self, x):
        return super().forward(x), None


class RMSNorm(nn.Module):
    """模拟原代码的RMSNorm归一化（PyTorch原生实现）"""

    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))  # 缩放权重
        self.eps = eps  # 避免除零的小值

    def forward(self, x):
        # RMSNorm公式：x * (x²的均值+eps)的平方根倒数 * 权重
        x_norm = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        output = x_norm * self.weight
        return output, None  # 适配原代码返回格式


def get_rope(rotary_dim, max_position=8192, base=10000, rope_scaling=None, is_neox_style=False):
    """模拟原代码的RoPE生成（Windows适配：移除NPU算子，用基础旋转逻辑）"""

    class RotaryEmbedding:
        def __init__(self):
            self.rotary_dim = rotary_dim  # RoPE作用的维度
            self.max_position = max_position  # 最大序列长度
            self.base = base  # RoPE基础频率
            self.rope_scaling = rope_scaling  # YARN缩放配置

        def get_cos_sin(self, positions):
            """生成positions对应的cos/sin矩阵（适配Windows CPU/GPU）"""
            seq_len = positions.max().item() + 1 if positions.dim() > 0 else 1
            # 计算频率：theta = base^(-2i/rotary_dim) (i为维度索引)
            theta = 1.0 / (self.base ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) / self.rotary_dim))

            # 应用YARN缩放（若配置）
            seq_idx = torch.arange(seq_len, dtype=torch.float32)
            if self.rope_scaling and self.rope_scaling.get("rope_type") == "deepseek_yarn":
                seq_idx = seq_idx / self.rope_scaling["factor"]  # 缩放序列位置

            # 频率矩阵：(seq_len, rotary_dim//2)
            freqs = torch.outer(seq_idx, theta)
            # 扩展维度：(1, seq_len, 1, rotary_dim)（适配Q/K的形状：(seq_len, num_heads, rotary_dim)）
            cos = freqs.cos().unsqueeze(0).unsqueeze(2).repeat_interleave(2, dim=-1)
            sin = freqs.sin().unsqueeze(0).unsqueeze(2).repeat_interleave(2, dim=-1)
            return cos, sin

    return RotaryEmbedding()


# 模拟原代码的model_extra_config（Windows单卡配置）
class MockModelExtraConfig:
    class OperatorOptConfig:
        merge_qkv = False  # 不合并QKV投影（简化逻辑）
        fa_quant = False  # 关闭量化（Windows环境暂不考虑）
        use_mlaprolog = False  # 关闭NPU Prolog算子
        prefill_enable_mla_alltoall = False  # 关闭分布式All-to-All
        enable_kv_rmsnorm_rope_cache = True  # 启用KV缓存复用
        decode_gear_list = [1, 2, 4]  # 模拟Decode阶段batch_size列表

    class ParallConfig:
        o_proj_tp_size = 1  # 输出投影无张量并行（单卡）
        dp_size = 1  # 无数据并行（单卡）

    operator_opt_config = OperatorOptConfig()
    parall_config = ParallConfig()


model_extra_config = MockModelExtraConfig()


# 模拟vllm的并行状态函数（Windows单卡返回1）
def get_tensor_model_parallel_world_size():
    return 1  # 单卡环境，张量并行进程数=1


def get_tensor_model_parallel_rank():
    return 0  # 单卡环境，当前进程 rank=0


def get_o_proj_tp_group():
    # 模拟输出投影的TP组（单卡）
    class MockGroup:
        world_size = 1
        rank_in_group = 0

    return MockGroup()


def get_o_proj_dp_group():
    # 模拟输出投影的DP组（单卡）
    class MockGroup:
        world_size = 1
        device_group = None

    return MockGroup()


def get_npu_device_count():
    return 1  # Windows无NPU，返回1（兼容代码逻辑）


# -------------------------- 核心DeepseekMLA类（Windows适配版） --------------------------
class DeepseekMLA_Windows(nn.Module):
    def __init__(
            self,
            config: PretrainedConfig,
            hidden_size: int = 4096,
            num_heads: int = 32,
            qk_nope_head_dim: int = 48,
            qk_rope_head_dim: int = 16,
            v_head_dim: int = 64,
            q_lora_rank: int = 64,
            kv_lora_rank: int = 64,
            rope_scaling: dict = None,
    ) -> None:
        super().__init__()
        # 基础参数
        self.hidden_size = hidden_size  # 模型隐藏层维度
        self.qk_nope_head_dim = qk_nope_head_dim  # Q/K中不应用RoPE的维度
        self.qk_rope_head_dim = qk_rope_head_dim  # Q/K中应用RoPE的维度
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim  # Q/K总头维度
        self.v_head_dim = v_head_dim  # V的头维度
        self.q_lora_rank = q_lora_rank  # Q的LoRA低秩维度
        self.kv_lora_rank = kv_lora_rank  # KV的LoRA低秩维度
        self.num_heads = num_heads  # 总注意力头数
        self.tp_size = get_tensor_model_parallel_world_size()  # 单卡=1
        self.num_local_heads = num_heads // self.tp_size  # 本地头数=总头数
        self.scale = self.qk_head_dim ** -0.5  # 注意力缩放因子（1/sqrt(d_k)）
        self.quant_symbol = False  # Windows环境关闭量化

        # -------------------------- 1. LoRA投影层初始化（Windows适配） --------------------------
        self.merge_qkv = model_extra_config.operator_opt_config.merge_qkv
        if self.q_lora_rank is not None:
            # Q的LoRA A投影（降秩）
            self.q_a_proj = ReplicatedLinear(self.hidden_size, self.q_lora_rank, bias=False)
            # KV的LoRA A投影（降秩+QK的RoPE维度）
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False
            )
            # Q的LoRA归一化
            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
            # Q的LoRA B投影（升秩：低秩→原头维度）
            self.q_b_proj = ColumnParallelLinear(
                self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False
            )
        else:
            # 无LoRA：直接Q投影
            self.q_proj = ColumnParallelLinear(
                self.hidden_size, self.num_heads * self.qk_head_dim, bias=False
            )

        # KV的LoRA归一化
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
        # KV的LoRA B投影（升秩：低秩→QK无RoPE维度+V维度）
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False
        )

        # -------------------------- 2. 输出投影层（Windows单卡适配） --------------------------
        self.o_proj = ReplicatedLinear(
            self.num_heads * self.v_head_dim, self.hidden_size, bias=False
        )

        # -------------------------- 3. RoPE初始化（Windows适配） --------------------------
        if rope_scaling:
            rope_scaling["rope_type"] = "deepseek_yarn"  # 标记YARN RoPE
        self.rotary_emb = get_rope(
            qk_rope_head_dim, max_position=8192, base=10000,
            rope_scaling=rope_scaling, is_neox_style=False
        )

        # 应用YARN缩放：调整注意力缩放因子
        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            mscale = yarn_get_mscale(rope_scaling["factor"], float(mscale_all_dim))
            self.scale *= mscale * mscale

        # -------------------------- 4. 其他初始化（Windows适配） --------------------------
        self.W_UK = None  # K的无RoPE部分权重（单卡环境延迟初始化）
        self.W_UV = None  # V的权重（单卡环境延迟初始化）
        self.is_init = True
        # 预分配Decode阶段的归一化结果缓存（适配静态形状）
        self.norm_res = {}
        for batch_size in model_extra_config.operator_opt_config.decode_gear_list:
            self.norm_res[batch_size] = torch.zeros(
                [batch_size * self.tp_size, self.q_lora_rank], dtype=torch.float32  # Windows用float32兼容CPU
            )

    def _apply_rope(self, x, cos, sin, positions):
        """Windows适配的RoPE应用（替换原代码的npu_interleave_rope）"""
        # x: (seq_len, num_heads, rotary_dim) → Q/K的RoPE部分
        # cos/sin: (1, max_seq_len, 1, rotary_dim) → 预生成的旋转矩阵
        seq_len, num_heads, rotary_dim = x.shape
        # 提取当前positions对应的cos/sin（避免冗余计算）
        cos = cos[:, positions, :, :rotary_dim].expand(-1, -1, num_heads, -1)  # (1, seq_len, num_heads, rotary_dim)
        sin = sin[:, positions, :, :rotary_dim].expand(-1, -1, num_heads, -1)

        # 分离实部和虚部（RoPE核心逻辑：将维度分为两半，模拟复数旋转）
        x1 = x[..., :rotary_dim // 2]  # 实部
        x2 = x[..., rotary_dim // 2:]  # 虚部
        # 旋转公式：cos*实部 - sin*虚部, cos*虚部 + sin*实部
        rotated_x1 = x1 * cos[0, :, :, :rotary_dim // 2] - x2 * sin[0, :, :, :rotary_dim // 2]
        rotated_x2 = x2 * cos[0, :, :, :rotary_dim // 2] + x1 * sin[0, :, :, :rotary_dim // 2]
        # 合并旋转后的维度
        rotated_x = torch.cat([rotated_x1, rotated_x2], dim=-1)
        return rotated_x

    def forward(self, hidden_states: torch.Tensor, positions: torch.Tensor):
        """Windows适配的前向传播（仅保留核心Prefill逻辑）"""
        # hidden_states: (batch_size*seq_len, hidden_size) → 输入隐藏态
        # positions: (seq_len,) → 每个token的位置索引

        # -------------------------- 1. Q的LoRA前向 --------------------------
        q_lowrank, _ = self.q_a_proj(hidden_states)  # (seq_len, q_lora_rank) → LoRA A投影
        q, _ = self.q_a_layernorm(q_lowrank)  # 归一化
        q, _ = self.q_b_proj(q)  # (seq_len, num_heads*qk_head_dim) → LoRA B投影
        q = q.view(-1, self.num_local_heads, self.qk_head_dim)  # (seq_len, num_heads, qk_head_dim) # 将经过 LoRA 恢复的扁平特征（[64, 2048]），拆分成 “序列长度 × 头数 × 每个头的维度” 的形状（[64, 32, 64]），为后续多头注意力的并行计算做准备。

        # -------------------------- 2. KV的LoRA前向 --------------------------
        kv, _ = self.kv_a_proj_with_mqa(hidden_states)  # (seq_len, kv_lora_rank + qk_rope_head_dim)
        # 拆分KV的LoRA低秩部分和QK的RoPE部分
        kv_a, k_pe_raw = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv_a, _ = self.kv_a_layernorm(kv_a)  # KV归一化

        # -------------------------- 3. 应用RoPE（Windows适配） --------------------------
        # 拆分Q的无RoPE部分和RoPE部分
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
        # 生成cos/sin矩阵
        cos, sin = self.rotary_emb.get_cos_sin(positions)
        # Q应用RoPE
        q_pe_rotated = self._apply_rope(q_pe, cos, sin, positions)
        # 合并Q的两部分
        q = torch.cat([q_nope, q_pe_rotated], dim=-1)

        # K应用RoPE（QK的RoPE部分共享）
        k_pe_rotated = self._apply_rope(
            k_pe_raw.view(-1, 1, self.qk_rope_head_dim),  # (seq_len, 1, qk_rope_head_dim)
            cos, sin, positions
        ).squeeze(1)  # (seq_len, qk_rope_head_dim)

        # -------------------------- 4. KV的B投影（升秩） --------------------------
        kv_b, _ = self.kv_b_proj(kv_a)  # (seq_len, num_heads*(qk_nope_head_dim + v_head_dim))
        kv_b = kv_b.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
        # 拆分K的无RoPE部分和V
        k_nope, v = torch.split(kv_b, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        # 合并K的无RoPE部分和RoPE部分
        k = torch.cat([k_nope, k_pe_rotated.unsqueeze(1).repeat(1, self.num_local_heads, 1)], dim=-1)

        # -------------------------- 5. 注意力计算（Windows适配：用einsum实现） --------------------------
        # QK匹配：(seq_len, num_heads, qk_dim) @ (seq_len, num_heads, qk_dim) → (num_heads, seq_len, seq_len)
        attn_scores = torch.einsum("shd,thd->hst", q, k) * self.scale
        # 因果掩码：下三角矩阵（避免未来token泄露）
        seq_len = q.shape[0]
        mask = torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool))
        attn_scores = attn_scores.masked_fill(~mask, float("-inf"))  # 掩码位置设为负无穷
        # Softmax归一化
        attn_weights = torch.softmax(attn_scores, dim=-1)
        # V加权：(num_heads, seq_len, seq_len) @ (seq_len, num_heads, v_dim) → (seq_len, num_heads, v_dim)
        attn_output = torch.einsum("hst,thd->shd", attn_weights, v)

        # -------------------------- 6. 输出投影 --------------------------
        attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)  # 展平：(seq_len, num_heads*v_dim)
        output, _ = self.o_proj(attn_output)  # (seq_len, hidden_size) → 最终输出

        return output


# -------------------------- Windows环境独立执行验证 --------------------------
if __name__ == "__main__":
    # 1. 配置初始化
    config = PretrainedConfig(rms_norm_eps=1e-6)  # 模拟transformers配置
    # YARN RoPE配置：序列长度缩放2倍（8192→16384）
    rope_scaling = {"factor": 2.0, "mscale_all_dim": True}

    # 2. 模块初始化（Windows单卡）
    mla = DeepseekMLA_Windows(
        config=config,
        hidden_size=4096,
        num_heads=32,
        qk_nope_head_dim=48,
        qk_rope_head_dim=16,
        v_head_dim=64,
        q_lora_rank=64,
        kv_lora_rank=64,
        rope_scaling=rope_scaling
    ).to(torch.float32)  # Windows CPU/GPU兼容float32

    # 3. 构造输入（Windows环境：小序列长度方便测试）
    seq_len = 64  # 序列长度
    hidden_states = torch.randn(seq_len, 4096, dtype=torch.float32)  # (seq_len, hidden_size)
    positions = torch.arange(seq_len, dtype=torch.int64)  # 每个token的位置：0~63

    # 4. 前向计算
    with torch.no_grad():  # 推理模式，禁用梯度
        output = mla(hidden_states, positions)

    # 5. 验证结果
    print("=== DeepseekMLA Windows前向验证 ===")
    print(f"输入形状: {hidden_states.shape} → (seq_len, hidden_size)")
    print(f"输出形状: {output.shape} → (seq_len, hidden_size)")  # 预期与输入维度一致
    print(f"输出均值: {output.mean().item():.4f}")  # 数值合理性验证（应为接近0的浮点数）
    print(f"输出标准差: {output.std().item():.4f}")  # 数值合理性验证（应为接近1的浮点数）
